package cn.junki.sparkbigmodeldemo.websocket;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import cn.hutool.json.JSONUtil;
import cn.junki.sparkbigmodeldemo.entity.dto.SparkBigModelChatDTO;
import cn.junki.sparkbigmodeldemo.entity.vo.JsonResult;
import cn.junki.sparkbigmodeldemo.service.SparkBigModelService;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Collections;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket服务
 *
 * @author Junki
 * @date 2023-09-21 09:06:06
 */
@Slf4j
@Data
@ServerEndpoint("/demo/spark/chat")
@Component
public class SparkBigModelChatWebSocketServer {

    /**
     * AiChatWebSocketServer对象集合
     */
    public static final ConcurrentHashMap<String, SparkBigModelChatWebSocketServer> webSocketMap = new ConcurrentHashMap<>();

    /**
     * 当前会话
     */
    private Session session;

    /**
     * 会话id
     */
    private String sessionId;

    /**
     * 当前用户id
     */
    private String userId;

    /**
     * 连接建立成功
     */
    @OnOpen
    public void onOpen(Session session) {
        // 获取token
        String token = session.getRequestParameterMap().getOrDefault("token", Collections.singletonList("")).get(0);

        // 当前会话
        this.session = session;

        // TODO 根据token获取用户id
        String userId = token;

        // token无效
        if (StrUtil.isEmpty(userId)) {
            this.sendMessage(JsonResult.fail(403, "token无效"));
            return;
        }
        this.userId = userId;

        // TODO 随机一个sessionId，可根据sessionId存储和查询每次对话记录
        this.sessionId = this.userId + '-' + IdUtil.fastSimpleUUID();

        // 加入集合中
        webSocketMap.put(this.userId, this);

        log.info("[AiChat]用户连接:" + this.userId + ",当前会话数为:" + webSocketMap.size());
    }

    /**
     * 连接关闭
     */
    @OnClose
    public void onClose() {
        if (null == this.userId) {
            return;
        }
        // 从集合中删除
        webSocketMap.remove(this.userId);
        // 置空当前session
        this.session = null;
        log.info("[AiChat]用户退出:" + this.userId + ",当前会话数为:" + webSocketMap.size());
    }

    /**
     * 收到客户端消息
     *
     * @param message 客户端发送过来的消息
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        log.info("[AiChat]收到客户端发来的消息:{}", message);

        // 封装消息
        SparkBigModelChatDTO sparkBigModelChatDTO;
        try {
            sparkBigModelChatDTO = JSONUtil.toBean(message, SparkBigModelChatDTO.class);
        } catch (Exception e) {
            this.sendMessage(JsonResult.fail(400, "消息格式错误"));
            log.error("[AiChat]客户端消息格式化错误：", e);
            return;
        }

        // 消息不能为空
        if (null == sparkBigModelChatDTO || CollUtil.isEmpty(sparkBigModelChatDTO.getMessageList())) {
            this.sendMessage(JsonResult.fail(400, "消息不能为空"));
            return;
        }

        // 调用星火大模型
        try {
            SpringUtil.getBean(SparkBigModelService.class).chat(sparkBigModelChatDTO, (chatVO) -> {
                if (chatVO.isSuccess()) {
                    this.sendMessage(JsonResult.success(chatVO));
                } else {
                    this.sendMessage(JsonResult.fail(500, "系统繁忙，请稍后再试"));
                }
            });
        } catch (Exception e) {
            this.sendMessage(JsonResult.fail(500, "系统繁忙，请稍后再试"));
            log.error("[AiChat]调用星火大模型错误：", e);
        }
    }

    /**
     * 连接错误
     */
    @OnError
    public void onError(Session session, Throwable error) {
        log.error("[AiChat]用户错误:" + this.userId + "：", error);
    }

    /**
     * 服务器主动推送封装消息
     */
    public void sendMessage(JsonResult jsonResult) {
        if (null == this.session) {
            return;
        }
        try {
            this.session.getBasicRemote().sendText(jsonResult.toJsonStr());
        } catch (Exception e) {
            log.error("[AiChat]消息推送失败，用户:" + this.userId);
        }
        if (!jsonResult.isSuccess()) {
            try {
                this.session.close();
            } catch (IOException e) {
                log.error("[AiChat]关闭会话失败，用户:" + this.userId);
            }
        }
    }

    /**
     * 服务器主动推送消息
     */
    private void sendText(String text) {
        try {
            this.session.getBasicRemote().sendText(text);
        } catch (IOException e) {
            log.error("[AiChat]消息推送失败，用户:" + this.userId);
        }
    }

}
